import matplotlib.pyplot as plt
import numpy as np
from numpy.typing import NDArray
from typing import Optional

from common_util import bounding_box_to_points, transform_3d_cloud_to_image, transform_3d_cloud

MATPLOTLIB_COLOR = {
    'blue': 'b',
    'green': 'g',
    'red': 'r',
    'cyan': 'c',
    'magenta': 'm',
    'yellow': 'y',
    'black': 'k',
    'white': 'w',
    'dimgrey': 'dimgrey',
    'grey': 'grey',
    'lightred': 'tomato',
    'orange': 'orange',
    'gold': 'gold',
    'lightgreen': 'lime',
    'lightblue': 'deepskyblue',
    'purple': 'darkviolet',
    'pink': 'hotpink'
}

def init_axes(figure:str=None, dpi: float=100) -> plt.Axes:
    fig = plt.figure(figure, dpi=dpi)
    return fig.add_axes([0, 0, 1, 1])

def matplot_2d_cloud(cloud_x, cloud_y, scale=0.1, rgb=MATPLOTLIB_COLOR['blue'],
                     alpha=1, label=None, axes: Optional[plt.Axes]=None) -> None:
    """ plot 2D cloud with matplotlib
    :param cloud_x: list of cloud x
    :param cloud_y: list of cloud y
    :param scale: size of cloud points, float or list of float
    :param rgb: color used for plot
    :param alpha: opacity of line, 0 for transparent, 1 for opaque
    :param label: label of cloud
    :param axes: axes handler of matplotlib
    """
    if axes is None:
        axes = init_axes(dpi=1000)
    if label is None:
        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none')
    else:
        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none', label=label)
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')

def matplot_2d_color_cloud(cloud_x, cloud_y, rgb ,scale=0.1,
                           alpha=1, label=None, axes: Optional[plt.Axes]=None) -> None:
    if axes is None:
        axes = init_axes(dpi=1000)
    if label is None:
        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none', cmap='jet')
    else:
        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none', label=label, cmap='jet')
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')

def matplot_cloud_to_image(cloud: NDArray, cloud_to_camera: NDArray, camera_intrinsic: NDArray,
                           crop=(-1, -1) , scale=0.5, rgb=MATPLOTLIB_COLOR['blue'], alpha=1, label=None,
                           axes: Optional[plt.Axes]=None) -> None:
    assert len(crop) == 2
    cloud = transform_3d_cloud_to_image(cloud, cloud_to_camera, camera_intrinsic)
    cloud = cloud[cloud[:, 2] > 0]
    if crop[0] > 0:
        cloud = cloud[(cloud[:, 0] >= 0) & (cloud[:, 0] <= crop[0])]
    if crop[1] > 0:
        cloud = cloud[(cloud[:, 1] >= 0) & (cloud[:, 1] <= crop[1])]
    if rgb == 'intensity':
        matplot_2d_color_cloud(cloud[:, 0], cloud[:, 1], cloud[:, -1], scale, alpha, label, axes)
    else:
        matplot_2d_cloud(cloud[:, 0], cloud[:, 1], scale, rgb, alpha, label, axes)

def matplot_image(np_image: NDArray, axes: Optional[plt.Axes]=None) -> None:
    if axes is None:
        axes = init_axes(dpi=300)
    axes.imshow(np_image)

def matplot_bounding_box(position: list, size: list, rotation: list, scale=1,
                         rgb=MATPLOTLIB_COLOR['green'], alpha=1, axes: Optional[plt.Axes]=None) -> None:
    """ plot 3D bounding box with matplotlib
    :param position: position of box
    :param size: size of box (x, y, z)
    :param rotation: heading of box in euler angles
    :param scale: line width of box
    :param rgb: color used for plot
    :param alpha: opacity of line, 0 for transparent, 1 for opaque
    :param axes: axes handler of matplotlib
    """
    if axes is None:
        axes = init_axes()
    box_pts = bounding_box_to_points(position, size, rotation)[:5, :2]
    axes.plot(box_pts[:, 0], box_pts[:, 1], c=rgb, lw=scale, alpha=alpha)
    axes.set_xlabel('x axis')
    axes.set_ylabel('y axis')
    axes.axis('equal')

def matplot_line(points_x, points_y, scale=0.1, rgb=MATPLOTLIB_COLOR['green'],
                 alpha=1, axes: Optional[plt.Axes]=None) -> None:
    if axes is None:
        axes = init_axes()
        axes.set_xlabel('x axis')
        axes.set_ylabel('y axis')
        axes.axis('equal')
    assert len(points_x) == len(points_y)
    axes.plot(points_x, points_y, c=rgb, lw=scale, alpha=alpha)

def matplot_line_to_image(points: NDArray, cloud_to_camera: NDArray, camera_intrinsic: NDArray,
                          crop=(-1, -1), scale=0.5, rgb=MATPLOTLIB_COLOR['blue'], alpha=1,
                          axes: Optional[plt.Axes]=None) -> None:
    points = transform_3d_cloud(points, cloud_to_camera[:3, 3], cloud_to_camera[:3, :3])
    if not np.all(points[:, 2] > 0) and points.shape[0] > 1:
        filled_pts = []
        for i in range(points.shape[0] - 1):
            filled_pts.append(points[i])
            if points[i, 2] * points[i+1, 2] <= 0:
                vector = points[i+1] - points[i]
                length = np.linalg.norm(vector)
                vector /= length
                for j in range(int(length)):
                    filled_pts.append(points[i] + j * vector)
        points = np.array(filled_pts)
    points = transform_3d_cloud_to_image(points, np.identity(4), camera_intrinsic)
    points = points[points[:, 2] > 0]
    x_flt, y_flt = np.array([True]*points.shape[0]), np.array([True]*points.shape[0])
    if crop[0] > 0:
        x_flt = (points[:, 0] > 0) & (points[:, 0] < crop[0])
    if crop[1] > 0:
        y_flt = (points[:, 1] > 0) & (points[:, 1] < crop[1])
    points = points[x_flt & y_flt]
    if points.shape[0] > 1:
        matplot_line(points[:, 0], points[:, 1], scale, rgb, alpha, axes)

def matplot_text(x: float, y: float, text: str, weight='light', scale='small',
                 rgb=MATPLOTLIB_COLOR['black'], alpha=1, axes: Optional[plt.Axes]=None) -> None:
    if axes is None:
        axes = init_axes()
    axes.text(x, y, text, weight=weight, size=scale, c=rgb, alpha=alpha)
